import os
import dify_helper
import pandas as pd
import argparse
import configparser
import json

def get_ocr_first(row:dict,query:str,dh:dify_helper.DifyHelper):
    paramters = {
        "ocr_content":row["OCR内容"]
    }
    try:
        answer = dh.send_json_message_get_json_result(paramters,query)
    except Exception as ex:
        print(ex)
        answer = "解析失败"
    # print(answer)
    return answer

def get_ocr_second(row:dict,query:str,condition_col:str,dh:dify_helper.DifyHelper):
    paramters = {
        "ocr_content":row["OCR内容"]
    }
    try:
        cc = json.loads(row[condition_col])
        if cc['is_clear'] == False or cc['report_type'] in ['无法确定','非医学影像报告']:
            return ""
        answer = dh.send_json_message(paramters,query)
    except Exception as ex:
        print(ex)
        answer = ""
    return answer

def get_ocr_third(row:dict,query:str,condition_col:str,dh:dify_helper.DifyHelper):
    paramters = {
        "ocr_content":row[condition_col],
        "hospital_name":'清华长庚医院'
    }
    try:
        if row[condition_col] == "":
            return ""
        answer = dh.send_json_message(paramters,query)
    except Exception as ex:
        print(ex)
        answer = "解读失败"
    return answer

def get_ocr_fourth(row:dict,query:str,condition_col:str,dh:dify_helper.DifyHelper):
    paramters = {
        "ocr_content":row[condition_col]
    }
    try:
        if row[condition_col] == "":
            return ""
        answer = dh.send_json_message_get_json_result(paramters,query)
    except Exception as ex:
        print(ex)
        answer = "解析失败"
    return answer

def run_dify_batch(data_excel_path:str,out_excel_path:str,dh_arr:list[dict]):
    data_df = pd.read_excel(data_excel_path)
    # for dh in dh_arr:
    #     data_df[dh['agent_name']] = data_df.apply(lambda x:run_dify(x,"开始解读",dh['dh']),axis=1)
    dh1 = dh_arr[0]
    data_df[dh1['agent_name']] = data_df.apply(lambda x:get_ocr_first(x,"开始",dh1['dh']),axis=1)
    # data_df["is_clear"] = data_df.apply(lambda x: json.loads(x[dh1['agent_name']])['is_clear'],axis=1)
    # data_df["report_type"] = data_df.apply(lambda x: json.loads(x[dh1['agent_name']])['report_type'],axis=1)

    dh2 = dh_arr[1]
    data_df[dh2['agent_name']] = data_df.apply(lambda x:get_ocr_second(x,"开始",dh1['agent_name'],dh2['dh']),axis=1)

    dh3 = dh_arr[2]
    data_df[dh3['agent_name']] = data_df.apply(lambda x:get_ocr_third(x,"开始解读",dh2['agent_name'],dh3['dh']),axis=1)

    dh4 = dh_arr[3]
    data_df[dh4['agent_name']] = data_df.apply(lambda x:get_ocr_fourth(x,"开始",dh2['agent_name'],dh4['dh']),axis=1)
    data_df.to_excel(out_excel_path)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="dify批量运行器")
    parser.add_argument('-rep',"--raw_excel_path", type = str, help = "原始excel文件地址")
    parser.add_argument('-oep',"--out_excel_path", type = str, help = "输出excel文件地址")
    args = parser.parse_args()
    config = configparser.ConfigParser()
    config.read("config.ini")
    print(config.sections())
    raw_excel_path = args.raw_excel_path
    out_excel_path = args.out_excel_path
    agent_name_arr_str = config['dify']['agent_name_arr']
    agent_name_arr = agent_name_arr_str.split("；")
    print(agent_name_arr)
    api_key_arr_str = config['dify']['api_key_arr']
    api_key_arr =  api_key_arr_str.split("；")
    print(api_key_arr)
    dh_arr = []
    for index,api_key in enumerate(api_key_arr):
        dh_arr.append({'agent_name':agent_name_arr[index],'dh':dify_helper.DifyHelper(config['dify']['base_url'],api_key)})
    run_dify_batch(raw_excel_path,out_excel_path,dh_arr)